import gymnasium as gym
from stable_baselines3 import PPO
from stable_baselines3.common.callbacks import BaseCallback
from stable_baselines3.common.vec_env import SubprocVecEnv
import torch
import time
from swimmer_pid_env import SwimmerPIDEnv  # 导入自定义的 Swimmer PID 环境

class TimeCallback(BaseCallback):
    def __init__(self, total_timesteps, verbose=0):
        super(TimeCallback, self).__init__(verbose)
        self.total_timesteps = total_timesteps
        self.start_time = None

    def _on_training_start(self):
        # 记录训练开始时间
        self.start_time = time.time()

    def _on_step(self):
        # 计算并打印训练经过的时间和步数
        elapsed_time = time.time() - self.start_time
        num_timesteps = self.num_timesteps
        print(f"Elapsed time: {elapsed_time:.2f}s, Timesteps: {num_timesteps}/{self.total_timesteps}")
        return True

class LoggingCallback(BaseCallback):
    def __init__(self, verbose=0):
        super(LoggingCallback, self).__init__(verbose)

    def _on_step(self):
        # 记录损失函数和其他信息
        if self.locals.get('infos'):
            for info in self.locals['infos']:
                if 'episode' in info.keys():
                    print(f"Episode reward: {info['episode']['r']}, length: {info['episode']['l']}")
        return True

def make_env():
    def _init():
        env = SwimmerPIDEnv()
        env.reset()  # 确保在调用 render 之前调用 reset
        return env
    return _init

# 主函数
def main():
    num_envs = 10  # 并行环境数量
    env = SubprocVecEnv([make_env() for _ in range(num_envs)])  # 创建并行环境
    total_timesteps = 6000000

    # 选择设备（GPU 或 CPU）
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # 创建 PPO 模型
    model = PPO('MlpPolicy', env, verbose=1, device=device)

    # 创建回调函数实例
    logging_callback = LoggingCallback()
    time_callback = TimeCallback(total_timesteps=total_timesteps)

    # 训练模型
    model.learn(total_timesteps=total_timesteps, callback=[time_callback, logging_callback])

    # 保存模型
    model.save("/home/sh/catkin_ws/src/ymbot_e_control/policy/ppo_swimmer_pid")

    # 关闭环境
    env.close()

if __name__ == '__main__':
    main()
